package com.david.kafka.auth.handler;

import com.alibaba.druid.pool.DruidDataSource;
import com.alibaba.fastjson.JSON;
import com.david.kafka.auth.util.DataSourceUtil;
import org.apache.kafka.common.security.JaasContext;
import org.apache.kafka.common.security.plain.PlainAuthenticateCallback;
import org.apache.kafka.common.security.plain.PlainLoginModule;
import org.apache.kafka.common.security.plain.internals.PlainServerCallbackHandler;
import org.apache.kafka.common.utils.Utils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.security.crypto.bcrypt.BCryptPasswordEncoder;

import javax.security.auth.callback.Callback;
import javax.security.auth.callback.NameCallback;
import javax.security.auth.callback.UnsupportedCallbackException;
import javax.security.auth.login.AppConfigurationEntry;
import java.io.IOException;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.List;
import java.util.Map;
import java.util.Objects;

import static com.david.kafka.auth.constants.Constants.*;

/**
 * kafka认证逻辑 sasl/plain二次开发
 *
 * @author David
 * @date 2022.06.30
 */
public class DbAuthenticationCallbackHandler extends PlainServerCallbackHandler {

    private static final Logger LOG = LoggerFactory.getLogger(DbAuthenticationCallbackHandler.class);

    private static final String JAAS_USER_PREFIX = "user_";

    private DruidDataSource dataSource = null;

    private List<AppConfigurationEntry> jaasConfigEntries;

    /**
     * 存储基本配置 用户名密码  idaas配置
     */
    private Map<String, ?> options;

    /**
     * db 验证开关
     */
    private boolean enableDbAuth;

    /**
     * resourceSystemId
     */
    private String resourceSystemId;

    /**
     * tenantId
     */
    private String tenantId;

    @Override
    public void handle(Callback[] callbacks) throws IOException, UnsupportedCallbackException {
        String username = null;
        for (Callback callback : callbacks) {
            if (callback instanceof NameCallback) {
                username = ((NameCallback) callback).getDefaultName();
            } else if (callback instanceof PlainAuthenticateCallback) {
                PlainAuthenticateCallback plainCallback = (PlainAuthenticateCallback) callback;
                boolean authenticated = authenticate(username, plainCallback.password());
                plainCallback.authenticated(authenticated);
                LOG.info("认证 username:{},result:{}", username, authenticated);
            } else {
                throw new UnsupportedCallbackException(callback);
            }
        }
    }

    @Override
    protected boolean authenticate(String username, char[] password) throws IOException {
        if (username == null || password == null) {
            LOG.error("用户名密码有空值!");
            return false;
        }

        //先配置文件读取
        try {
            boolean fileAuthenticateResult = fileAuthenticate(username, password);
            LOG.info("fileAuthenticateResult username:{},result:{}", username, fileAuthenticateResult);
            if (fileAuthenticateResult) {
                return Boolean.TRUE;
            }
        } catch (IOException e) {
            LOG.error("从配置文件读取验证失败{}", e);
            throw e;
        }

        //数据库验证
        if (!enableDbAuth) {
            return false;
        }
        return dbAuthenticate(username, password);
    }


    protected boolean fileAuthenticate(String username, char[] password) throws IOException {
        if (username == null) {
            return false;
        } else {
            String expectedPassword = JaasContext.configEntryOption(this.jaasConfigEntries, JAAS_USER_PREFIX + username, PlainLoginModule.class.getName());
            return expectedPassword != null && Utils.isEqualConstantTime(password, expectedPassword.toCharArray());
        }
    }

    protected boolean dbAuthenticate(String username, char[] passwordCharArray) throws IOException {
        LOG.info("begin dbAuthenticateResult username:{}", username);
        String password = new String(passwordCharArray);

        String userQuery = "SELECT\n" +
                "\t        t1.id,\n" +
                "\t        t1.username,\n" +
                "\t        t1.password,\n" +
                "            itu.tenant_id tenantId\n" +
                "        FROM\n" +
                "            idaas_user t1\n" +
                "            LEFT JOIN idaas_tenant_user itu on t1.id = itu.user_id\n" +
                "       WHERE \n" +
                "            itu.tenant_id = ? \n" +
                "            and t1.del_flag = 0\n" +
                "            AND t1.username = ?\n" +
                "            AND t1.status ='ENABLE'\n" +
                "        LIMIT 1";
        try (Connection conn = dataSource.getConnection();
             PreparedStatement statement = conn.prepareStatement(userQuery);
        ) {
            statement.setString(1, tenantId);
            statement.setString(2, username);

            try (ResultSet resultSet = statement.executeQuery()) {
                if (resultSet.next()) {
                    String dbPassword = resultSet.getString("password");
                    BCryptPasswordEncoder passwordEncoder = new BCryptPasswordEncoder();
                    Boolean bl = passwordEncoder.matches(password, dbPassword);
                    if (!bl) {
                        LOG.info("{}密码验证失败", username);
                        return bl;
                    }
                    LOG.info("{}密码验证通过", username);
                    //验证是否有系统权限
                    String userId = resultSet.getString("id");
                    Boolean roleAuthBoolean = validateRoleAuth(userId);
                    LOG.info("{}角色权限验证{}", username, roleAuthBoolean);
                    return roleAuthBoolean;
                } else {
                    LOG.info("认证失败,user {} not found in db.", username);
                    return false;
                }
            }
        } catch (SQLException e) {
            LOG.error("数据库查询用户异常{}", e);
            throw new RuntimeException(e);
        }
    }

    /**
     * 验证当前用户是否有系统的角色
     *
     * @return
     */
    private Boolean validateRoleAuth(String userId) {
        String sql = "SELECT\n" +
                "roleUser.user_id,\n" +
                "u.username,\n" +
                "roleUser.role_id,\n" +
                "rol.role_name,\n" +
                "rol.role_key\n" +
                "FROM\n" +
                "idaas_role_user AS roleUser\n" +
                "LEFT JOIN idaas_role rol ON\n" +
                "\trol.id = roleUser.role_id\n" +
                "LEFT JOIN idaas_user AS u ON\n" +
                "roleUser.user_id = u.id\n" +
                "WHERE\n" +
                "roleUser.user_id = ?\n" +
                "AND rol.resource_system_id = ?";

        try (Connection conn = dataSource.getConnection();
             PreparedStatement statement = conn.prepareStatement(sql);
        ) {
            statement.setString(1, userId);
            statement.setString(2, resourceSystemId);
            try (ResultSet resultSet = statement.executeQuery()) {
                if (resultSet.next()) {
                    return Boolean.TRUE;
                } else {
                    return Boolean.FALSE;
                }
            }
        } catch (SQLException e) {
            LOG.error("数据库查询角色资源异常{}", e);
            throw new RuntimeException(e);
        }
    }


    @Override
    public void configure(Map<String, ?> configs, String saslMechanism, List<AppConfigurationEntry> jaasConfigEntries) {
        //配置文件里面的用户名密码
        this.jaasConfigEntries = jaasConfigEntries;

        LOG.info("configs:{}", JSON.toJSONString(configs));

        String loginModuleName = PlainLoginModule.class.getName();
        LOG.info("loginModuleName:{}", loginModuleName);
        for (AppConfigurationEntry entry : jaasConfigEntries) {
            if (!loginModuleName.equals(entry.getLoginModuleName())) {
                continue;
            }
            options = entry.getOptions();
        }

        LOG.info("options:{}", JSON.toJSONString(options.toString()));

        if (options == null) {
            throw new RuntimeException(PlainLoginModule.class.getName() + " 缺少必须配置!");
        }


        Object endbAuthObject = configs.get(ENABLE_DB_AUTH);
        if (Objects.isNull(endbAuthObject)) {
            LOG.error("缺少配置enable_db_auth!");
            enableDbAuth = Boolean.FALSE;
            return;
        }

        enableDbAuth = TRUE.equalsIgnoreCase(endbAuthObject.toString());
        if (!enableDbAuth) {
            return;
        }

        resourceSystemId = (String) configs.get(IDAAS_RESOURCE_SYSTEMID);
        tenantId = (String) configs.get(IDAAS_TENANT_ID);

        dataSource = DataSourceUtil.getInstance(configs);
    }


    @Override
    public void close() {
        LOG.info("close");
        if (dataSource == null) {
            return;
        }
        dataSource.close();
    }

}
